package com.shujia.mllib

import org.apache.spark.ml.classification.{NaiveBayes, NaiveBayesModel}
import org.apache.spark.ml.evaluation.MulticlassClassificationEvaluator
import org.apache.spark.ml.feature.{HashingTF, IDF, IDFModel, Tokenizer}
import org.apache.spark.sql.{DataFrame, Dataset, Row, SparkSession}

object Demo09BayesTrain {
  def main(args: Array[String]): Unit = {
    val spark: SparkSession = SparkSession
      .builder()
      .master("local[*]")
      .appName("Demo07KMeans")
      .getOrCreate()

    import spark.implicits._

    val sentenceDF: DataFrame = spark
      .sparkContext
      .textFile("spark/data/mllib/data/bayesTrain.txt")
      .map(line => {
        val splits: Array[String] = line.split("\t")
        val label: Double = splits(0).toDouble
        // 使用IK分词器进行分词

        (label, splits(1), Demo08IK.fit(splits(1)))
      })
      .filter(_._3 != "")
      .toDF("label", "text", "sentence")
      .repartition(16)

    // 英文分词器：可以将数据按照空格或者是标点进行分词
    val tokenizer: Tokenizer = new Tokenizer().setInputCol("sentence").setOutputCol("words")
    val wordsData: DataFrame = tokenizer.transform(sentenceDF)
    wordsData.show()

    // TF模型：计算词语出现的频率
    val hashingTF: HashingTF = new HashingTF()
      .setInputCol("words")
      .setOutputCol("rawFeatures")
      .setNumFeatures(Math.pow(2, 18).toInt) // 数据集经过分词之后有多少个词语 就可以设置多少个特征 默认2的18次方

    val featurizedData: DataFrame = hashingTF.transform(wordsData)

    featurizedData.show(truncate = false)
    // alternatively, CountVectorizer can also be used to get term frequency vectors

    // IDF模型
    val idf: IDF = new IDF().setInputCol("rawFeatures").setOutputCol("features")
    val idfModel: IDFModel = idf.fit(featurizedData)

    val rescaledData: DataFrame = idfModel.transform(featurizedData)
    rescaledData.show(truncate = false)

    val dataArr: Array[Dataset[Row]] = rescaledData.randomSplit(Array(0.8, 0.2))
    val trainDF: Dataset[Row] = dataArr(0)
    val testDF: Dataset[Row] = dataArr(1)

    // Train a NaiveBayes model.
    val bayesModel: NaiveBayesModel = new NaiveBayes()
      .fit(trainDF)

    // Select example rows to display.
    val predictions: DataFrame = bayesModel.transform(testDF)
    predictions.show()

    // Select (prediction, true label) and compute test error
    val evaluator: MulticlassClassificationEvaluator = new MulticlassClassificationEvaluator()
      .setLabelCol("label")
      .setPredictionCol("prediction")
      .setMetricName("accuracy")
    val accuracy: Double = evaluator.evaluate(predictions)
    println(s"Test set accuracy = $accuracy")

    // 保存模型
    idfModel.write.overwrite().save("spark/data/mllib/bayes/idf")
    bayesModel.write.overwrite().save("spark/data/mllib/bayes/bModel")


  }

}
